#!/usr/bin/env python
import re
import glob
import time
import Queue
import socket
import threading
import xml.etree.ElementTree as etree


import libvirt

from utils import flatten
from xmlbuilder import XMLBuilder
from cmd_executor import SSHCMDExecutor

from network import get_next_mac, get_ips_for_hws
from devices import Device, ETHDevice
from storage import Storage

storage = Storage()

class virDomainEx(libvirt.virDomain):
    
    if_re = re.compile("inet addr:(?P<ip>(?:\d{1,3}\.){3}\d{1,3})")
    
    def __init__(self, *dt, **mp):
        # faked code for pylint only
        libvirt.virDomain.__init__(self, *dt, **mp)
        self.etree = None
        self.memory = None
        self.vcpu = None
        self.devices = None
        self.name = None
        self.persistent = None
        self.dom_type = None
        
    def save_network(self):
        return storage.set_vm(self.name,
                       dict((eth.hw, eth.ip)
                            for eth in self.eths(with_ip = True)))
    
    @classmethod
    def construct(cls, conn, persistent, name, memory, vcpu, *devices):
        etree = cls.makeXML(name, memory, vcpu, *devices)
        
        self = conn.createXML(str(etree), 0)
        self.__class__ = cls
        
        self.etree = etree
        self.memory = memory
        self.vcpu = vcpu
        self.devices = devices
        self.name = name
        self.persistent = persistent
        
        return self
    
    def update_net_from_storage(self):
        try:
            hwip = storage.get_vm(self.name)
            
            hw2dev = dict((eth.hw, eth) for eth in self.eths())
            
            hws = set(hwip.keys()).intersection(set(hw2dev.keys()))
            
            for hw in hws:
                hw2dev[hw].ip = hwip[hw]
            
        except KeyError:
            pass
    
    @classmethod
    def fromVirDomain(cls, domain):
        domain.__class__ = cls
        domain.etree = etree.fromstring(domain.XMLDesc(0))
        
        
        find = domain.etree.find
        findall = domain.etree.findall
        find_text = lambda path : find(path).text.strip()
        
        domain.memory = int(find_text('memory'))
        domain.vcpu = int(find_text('vcpu'))
        domain.dom_type = 'xen'
        
        domain.devices = []

        for obj in findall("devices/interface[@type='network']"):
            domain.devices.append(
                ETHDevice(
                    obj.find('mac').attrib['address'],
                    obj.find('target').attrib['dev'],
                    obj.find('source').attrib['network']))

        domain.name = find_text('name')
        domain.persistent = True
        
        domain.update_net_from_storage()
           
        return domain
    
    def toXML(self):
        return self.makeXML(self.name, self.memory, self.vcpu, self.devices)
    
    @classmethod
    def makeXML(cls, name, memory, vcpu, *devices):
        root = XMLBuilder(format=True)

        with root('domain', type='kvm'):
            root << ('name', name)
            root << ('memory', str(memory))
            root.uuid
            root << ('vcpu', str(vcpu))
            
            with root('os'):
                root << ('type', 'hvm')
                root << ('loader', '/usr/lib/xen-default/boot/hvmloader')
                root << ('boot', {'dev':'hd'})
                root << ('boot', {'dev':'cdrom'})
                root << ('bootmenu', {'enable':'yes'})
                root << ('bios', {'useserial':'yes'})
            
            root << ('clock', {'sync':'localtime'})
            root << ('on_poweroff', 'destroy')
            
            with root.features:
                root.acpi
                root.hap
                root.apic
                
            with root('devices'):
                #root << ('emulator', '/usr/lib/xen-default/bin/qemu-dm')
                root << ('emulator', '/usr/bin/kvm')
                
                for dev in devices: 
                    dev.toxml(root)
                    
                with root('serial', type='pty'):
                    root << ('target', {'port':'0'})
                
                with root('console', type='pty'):
                    root << ('target', {'port':'0'})
                
                root << ('input', {'bus':'ps2', 'type':'mouse'})
                root << ('graphics', {'autoport':'yes', 
                                      'keymap':'en-us', 
                                      'type':'vnc', 
                                      'port':'-1'})
        
        return root
    
    def get_devices(self, tp = Device):
        for dev in self.devices:
            if isinstance(dev, tp):
                yield dev
    
    def __enter__(self):
        return self
    
    def __exit__(self, x, y, z):
        if not self.persistent:
            self.destroy()

    def get_ips(self):
        return [eth.ip for eth in self.eths(True) if eth.ip is not None]

    def eths(self, with_ip=False, conn=None, nemask=None):
        devs = list(self.get_devices(ETHDevice))

        if with_ip:
            no_ips = [dev for dev in devs if dev.ip is None]
            
            if no_ips:
                if nemask:
                    hw2ip = get_ips_for_hws(netmask,
                                            *[dev.hw for dev in no_ips])
                elif conn:
                    nmappings = {}
                    hw2ip = {}
                    
                    for dev in no_ips:
                        nmappings.setdefault(dev.network,[]).append(dev)
                
                    for network, curr_devs in nmappings.items():
                        net_xml = conn.networkLookupByName(network).XMLDesc()
                        tag = etree.fromstring(net_xml).find('ip')
                        
                        addr = tag.attrib['address']
                        mask = tag.attrib['netmask']
                        
                        vir_netmask = str(ipaddr.IPNetwork(
                                        "{0}/{1}".format(addr, mask)))
                        
                        hw2ip.update(get_ips_for_hws(vir_netmask,
                                                *[dev.hw for dev in curr_devs]))
                
                for dev in no_ips:
                    dev.ip = hw2ip.get(dev.hw, None)

        return devs

class VMConnector(object):
    def __init__(self, vm):
        self.ssh_addr = None
        self.rpyc_addr = None
        self.rpyc_pem_file = None        
        self.ssh_user = None
        self.ssh_passwd = None
        self.vm = vm

    def set_ssh_credentials(self, user, passwd):
        self.ssh_user = user
        self.ssh_passwd = passwd
        
    def check_ips(self, func, cache_attr = None):
        
        if cache_attr is not None:
            val = getattr(self, cache_attr)
        else:
            val = None

        ips = set(self.vm.get_ips())
        
        if val is not None and val in ips:
            ips.remove(val)
            search_ips = (val,) + tuple(ips)
        else:
            search_ips = ips
        
        for ip in search_ips:
            try:
                res = func(ip)
                if cache_attr is not None:
                    setattr(self, cache_attr, ip)
                return res
            except socket.timeout:
                pass
            
    def wait_ssh_ready(self):
        while True:
            res = self.conn_ssh()

            if res is not None:
                return res

            time.sleep(1)
    
    def conn_ssh(self):
        def conn_ssh(ip):
            return SSHCMDExecutor(ip,
                                  self.ssh_user,
                                  self.ssh_passwd,
                                  timeout = 10)
        
        res = self.check_ips(conn_ssh, 'ssh_addr')
        return res
    
    def set_rpyc_pem_file(self, pem_file):
        self.rpyc_pem_file = pem_file
    
    def rpyc_start_server(self,
                          pem_file_path = '/opt/rpyc/cert.pem',
                          rpyc_server = '/opt/rpyc/rpyc_classic.py'):
        ssh_conn = self.conn_ssh()
        ssh_conn.exec_simple_check(("nohup python {0} --ssl-keyfile " + \
                                   "{1} --ssl-certfile {1} &").format(
                                        rpyc_server,
                                        pem_file_path))
        
    def conn_rpyc(self):
        import rpyc
        
        def conn_rpyc(ip):
            return rpyc.classic.ssl_connect(
                                    ip,
                                    keyfile = self.rpyc_pem_file,
                                    certfile = self.rpyc_pem_file)
        
        return self.check_ips(conn_rpyc, 'rpyc_addr')

    
class virConnectEx(libvirt.virConnect):
    def lookupByName(self, name):
        return virDomainEx.fromVirDomain(
            libvirt.virConnect.lookupByName(self, name))
        
    def lookupByID(self, did):
        return virDomainEx.fromVirDomain(
            libvirt.virConnect.lookupByID(self, did))
        
    def allDomains(self, name_filter = None):
        if name_filter is not None:
            nfl = lambda name : glob.fnmatch.fnmatchcase(name, name_filter)
        else:
            nfl = lambda name : True
        
        for domain_id in self.listDomainsID():
            dom = self.lookupByID(domain_id)
            if nfl(dom.name):
                yield dom
    
    def networkDev(self, conn, network):
        netxml = conn.networkLookupByName(network).XMLDesc(0)
        return etree.fromstring(netxml).find('bridge').attrib['name']
        
    def allMacks(self):
        return flatten(hwip.keys() for _, hwip in storage.get_all())

    def allIPs(self):
        return flatten(hwip.values() for _, hwip in storage.get_all())

    def findFreeMac(self):
        used_macs = self.allMacks()

        for mac in get_next_mac():
            if mac not in used_macs:
                return mac


def open_libvirt(*dt, **mp):
    conn = libvirt.open(*dt, **mp)

    if conn.__class__ is not virConnectEx:
        conn.__class__ = virConnectEx
    
    conn.dom_type = 'kvm'

    return conn

